import tensorflow as tf


def top_k_accuracy(output, target, top_k=(1,)):
    max_k = max(top_k)
    batch_size = target.shape[0]

    prediction = tf.math.top_k(output, max_k).indices
    prediction = tf.transpose(prediction, perm=[1, 0])
    target_bd = tf.broadcast_to(target, prediction.shape)
    correct = tf.equal(prediction, target_bd)

    result = []
    for k in top_k:
        correct_k = tf.cast(correct[:k], dtype=tf.float32)
        correct_k = tf.reduce_sum(correct_k)
        accuracy = float(correct_k / batch_size)
        result.append(accuracy)

    return result


output = tf.random.normal([10, 6])
output = tf.math.softmax(output, axis=1)
target = tf.random.uniform([10], maxval=6, dtype=tf.int32)
prediction = tf.argmax(output, axis=1)

print('probability:\n', output.numpy())
print('probability top 6:\n', tf.math.top_k(output, 6).indices.numpy())
print('prediction:', prediction.numpy())
print('label:', target.numpy())

tk_accuracy = top_k_accuracy(output, target, top_k=(1, 2, 3, 4, 5, 6))
print('top accuracy:\n', tk_accuracy)
